import torch
import os
import glob
from PIL import Image
from torchvision.utils import save_image
from os.path import isfile
from torch.autograd import Variable
from utils.utils import *
from options.test_options import TestOptions
from models.transformer_arch import TransformerNet

def main():
    opt = TestOptions().parse()  # 解析模型测试命令行参数
    workspace_dir = make_dirs(f"{opt.output_dir}/{opt.name}")  # 创建工作空间文件夹
    device = torch.device(f"cuda:{opt.gpu_id}" if opt.gpu_id >= 0 else "cpu")  # 初始化测试设备

    # 初始化网络
    transform = style_image_transform()
    transformer_net = TransformerNet().to(device)
    transformer_net.load_state_dict(torch.load(opt.model))
    transformer_net.eval()

    # 进行风格迁移
    content_images = [opt.content] if isfile(opt.content) else glob.glob(f"{opt.content}{os.sep}*")
    for content_path in content_images:
        image = Variable(transform(Image.open(content_path))).to(device)
        image = image.unsqueeze(0)
        with torch.no_grad():
            stylized_image = denormalize_image(transformer_net(image)).cpu()

        filename = parse_filename(content_path)
        save_image(stylized_image, f"{workspace_dir}/{filename}.jpg")

        if opt.comparison:
            comparison_image = torch.cat((denormalize_image(image).cpu(), stylized_image), 2)
            save_image(comparison_image, f"{workspace_dir}/{filename}_comparison.jpg")

        print(f"{content_path} is completed.")

if __name__ == "__main__":
    main()
